#!/usr/bin/env python3
"""
Database manager for AI-Powered SSH Terminal
SQLite-based storage for connections, commands, and analytics
"""

import sqlite3
import json
import logging
from pathlib import Path
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
from contextlib import contextmanager

from models import ConnectionInfo, CommandHistory, ServerInfo, AISuggestion, AnalyticsData

logger = logging.getLogger(__name__)

class DatabaseManager:
    """SQLite database manager for terminal data"""
    
    def __init__(self, db_path: str = "ai_terminal.db"):
        self.db_path = Path(db_path)
        self.init_database()
    
    def init_database(self):
        """Initialize database with required tables"""
        with self.get_connection() as conn:
            # Connections table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS connections (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    hostname TEXT NOT NULL,
                    port INTEGER DEFAULT 22,
                    username TEXT NOT NULL,
                    enterprise TEXT,
                    tags TEXT,  -- JSON array
                    description TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_used TIMESTAMP,
                    usage_count INTEGER DEFAULT 0,
                    UNIQUE(hostname, port, username)
                )
            """)
            
            # Command history table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS command_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    connection_id INTEGER,
                    command TEXT NOT NULL,
                    command_type TEXT DEFAULT 'linux',
                    success BOOLEAN DEFAULT TRUE,
                    output TEXT,
                    execution_time REAL,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (connection_id) REFERENCES connections (id)
                )
            """)
            
            # Server information table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS server_info (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    connection_id INTEGER,
                    os_type TEXT,
                    os_version TEXT,
                    cpu_info TEXT,
                    memory_info TEXT,
                    disk_info TEXT,
                    network_info TEXT,
                    services TEXT,  -- JSON array
                    last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (connection_id) REFERENCES connections (id),
                    UNIQUE(connection_id)
                )
            """)
            
            # AI suggestions table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS ai_suggestions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    query TEXT NOT NULL,
                    suggested_command TEXT NOT NULL,
                    command_type TEXT DEFAULT 'linux',
                    success_rate REAL DEFAULT 0.0,
                    usage_count INTEGER DEFAULT 0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(query, command_type)
                )
            """)
            
            # Enterprise analytics table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS enterprise_analytics (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    enterprise TEXT NOT NULL,
                    metric_name TEXT NOT NULL,
                    metric_value TEXT,  -- JSON data
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(enterprise, metric_name)
                )
            """)
            
            conn.commit()
            logger.info("Database initialized successfully")
    
    @contextmanager
    def get_connection(self):
        """Get database connection with context manager"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
        except Exception as e:
            conn.rollback()
            logger.error(f"Database error: {e}")
            raise
        finally:
            conn.close()
    
    def store_connection(self, connection: ConnectionInfo) -> int:
        """Store connection information"""
        with self.get_connection() as conn:
            cursor = conn.execute("""
                INSERT OR REPLACE INTO connections 
                (name, hostname, port, username, enterprise, tags, description)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (
                connection.name,
                connection.hostname,
                connection.port,
                connection.username,
                connection.enterprise,
                json.dumps(connection.tags),
                connection.description
            ))
            connection_id = cursor.lastrowid
            conn.commit()
            logger.info(f"Stored connection: {connection.name}")
            return connection_id
    
    def get_connections(self, enterprise: Optional[str] = None) -> List[Dict]:
        """Get all connections, optionally filtered by enterprise"""
        with self.get_connection() as conn:
            if enterprise:
                cursor = conn.execute("""
                    SELECT * FROM connections 
                    WHERE enterprise = ? 
                    ORDER BY last_used DESC, name ASC
                """, (enterprise,))
            else:
                cursor = conn.execute("""
                    SELECT * FROM connections 
                    ORDER BY last_used DESC, name ASC
                """)
            
            connections = []
            for row in cursor.fetchall():
                connection = dict(row)
                connection['tags'] = json.loads(connection['tags'] or '[]')
                connections.append(connection)
            
            return connections
    
    def update_connection_usage(self, connection_id: int):
        """Update connection usage statistics"""
        with self.get_connection() as conn:
            conn.execute("""
                UPDATE connections 
                SET last_used = CURRENT_TIMESTAMP, 
                    usage_count = usage_count + 1
                WHERE id = ?
            """, (connection_id,))
            conn.commit()
    
    def store_command_history(self, history: CommandHistory) -> int:
        """Store command execution history"""
        with self.get_connection() as conn:
            cursor = conn.execute("""
                INSERT INTO command_history 
                (connection_id, command, command_type, success, output, execution_time)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (
                history.connection_id,
                history.command,
                history.command_type,
                history.success,
                history.output,
                history.execution_time
            ))
            history_id = cursor.lastrowid
            conn.commit()
            return history_id
    
    def get_command_history(self, connection_id: Optional[int] = None, limit: int = 100) -> List[Dict]:
        """Get command history"""
        with self.get_connection() as conn:
            if connection_id:
                cursor = conn.execute("""
                    SELECT * FROM command_history 
                    WHERE connection_id = ? 
                    ORDER BY timestamp DESC 
                    LIMIT ?
                """, (connection_id, limit))
            else:
                cursor = conn.execute("""
                    SELECT ch.*, c.name as connection_name 
                    FROM command_history ch
                    LEFT JOIN connections c ON ch.connection_id = c.id
                    ORDER BY ch.timestamp DESC 
                    LIMIT ?
                """, (limit,))
            
            return [dict(row) for row in cursor.fetchall()]
    
    def store_server_info(self, server_info: ServerInfo):
        """Store server information"""
        with self.get_connection() as conn:
            conn.execute("""
                INSERT OR REPLACE INTO server_info 
                (connection_id, os_type, os_version, cpu_info, memory_info, 
                 disk_info, network_info, services)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                server_info.connection_id,
                server_info.os_type,
                server_info.os_version,
                server_info.cpu_info,
                server_info.memory_info,
                server_info.disk_info,
                server_info.network_info,
                json.dumps(server_info.services or [])
            ))
            conn.commit()
    
    def get_server_info(self, connection_id: int) -> Optional[Dict]:
        """Get server information"""
        with self.get_connection() as conn:
            cursor = conn.execute("""
                SELECT * FROM server_info WHERE connection_id = ?
            """, (connection_id,))
            row = cursor.fetchone()
            
            if row:
                server_info = dict(row)
                server_info['services'] = json.loads(server_info['services'] or '[]')
                return server_info
            return None
    
    def store_ai_suggestion(self, suggestion: AISuggestion):
        """Store or update AI suggestion"""
        with self.get_connection() as conn:
            conn.execute("""
                INSERT OR REPLACE INTO ai_suggestions 
                (query, suggested_command, command_type, success_rate, usage_count)
                VALUES (?, ?, ?, ?, ?)
            """, (
                suggestion.query,
                suggestion.suggested_command,
                suggestion.command_type,
                suggestion.success_rate,
                suggestion.usage_count
            ))
            conn.commit()
    
    def get_ai_suggestions(self, query: str, command_type: str) -> List[Dict]:
        """Get AI suggestions for a query"""
        with self.get_connection() as conn:
            cursor = conn.execute("""
                SELECT * FROM ai_suggestions 
                WHERE query LIKE ? AND command_type = ?
                ORDER BY success_rate DESC, usage_count DESC
                LIMIT 5
            """, (f"%{query}%", command_type))
            
            return [dict(row) for row in cursor.fetchall()]
    
    def update_suggestion_feedback(self, suggestion_id: int, success: bool):
        """Update suggestion success rate based on feedback"""
        with self.get_connection() as conn:
            # Get current stats
            cursor = conn.execute("""
                SELECT success_rate, usage_count FROM ai_suggestions WHERE id = ?
            """, (suggestion_id,))
            row = cursor.fetchone()
            
            if row:
                current_rate = row['success_rate']
                usage_count = row['usage_count']
                
                # Calculate new success rate
                total_successes = current_rate * usage_count
                if success:
                    total_successes += 1
                usage_count += 1
                new_rate = total_successes / usage_count
                
                conn.execute("""
                    UPDATE ai_suggestions 
                    SET success_rate = ?, usage_count = ?, updated_at = CURRENT_TIMESTAMP
                    WHERE id = ?
                """, (new_rate, usage_count, suggestion_id))
                conn.commit()
    
    def get_analytics(self) -> AnalyticsData:
        """Get comprehensive analytics data"""
        with self.get_connection() as conn:
            # Total connections
            total_connections = conn.execute("SELECT COUNT(*) FROM connections").fetchone()[0]
            
            # Total commands
            total_commands = conn.execute("SELECT COUNT(*) FROM command_history").fetchone()[0]
            
            # Success rate
            success_count = conn.execute("""
                SELECT COUNT(*) FROM command_history WHERE success = 1
            """).fetchone()[0]
            success_rate = (success_count / total_commands * 100) if total_commands > 0 else 0
            
            # Top commands
            top_commands_cursor = conn.execute("""
                SELECT command, COUNT(*) as count, 
                       AVG(CASE WHEN success THEN 1.0 ELSE 0.0 END) as success_rate
                FROM command_history 
                GROUP BY command 
                ORDER BY count DESC 
                LIMIT 10
            """)
            top_commands = [dict(row) for row in top_commands_cursor.fetchall()]
            
            # Enterprise usage with proper field mapping
            enterprise_cursor = conn.execute("""
                SELECT enterprise as name, COUNT(*) as connection_count,
                       COALESCE(MAX(last_used), datetime('now')) as last_activity
                FROM connections 
                WHERE enterprise IS NOT NULL
                GROUP BY enterprise
                ORDER BY connection_count DESC
            """)
            
            enterprise_data = []
            for row in enterprise_cursor.fetchall():
                enterprise_dict = dict(row)
                
                # Get common commands for this enterprise
                cmd_cursor = conn.execute("""
                    SELECT ch.command, COUNT(*) as count
                    FROM command_history ch
                    JOIN connections c ON ch.connection_id = c.id
                    WHERE c.enterprise = ?
                    GROUP BY ch.command
                    ORDER BY count DESC
                    LIMIT 5
                """, (enterprise_dict['name'],))
                
                common_commands = [cmd_row['command'] for cmd_row in cmd_cursor.fetchall()]
                
                # Set required fields for EnterpriseInfo model
                enterprise_dict['common_commands'] = common_commands or ['No commands yet']
                enterprise_dict['server_types'] = ['Linux']  # Default server type
                
                # Convert last_activity to datetime if it's a string
                if isinstance(enterprise_dict['last_activity'], str):
                    try:
                        enterprise_dict['last_activity'] = datetime.fromisoformat(enterprise_dict['last_activity'].replace('Z', '+00:00'))
                    except ValueError:
                        enterprise_dict['last_activity'] = datetime.now()
                elif enterprise_dict['last_activity'] is None:
                    enterprise_dict['last_activity'] = datetime.now()
                
                enterprise_data.append(enterprise_dict)
            
            # Recent activity
            recent_cursor = conn.execute("""
                SELECT ch.command, ch.timestamp, c.name as connection_name
                FROM command_history ch
                LEFT JOIN connections c ON ch.connection_id = c.id
                ORDER BY ch.timestamp DESC
                LIMIT 20
            """)
            recent_activity = [dict(row) for row in recent_cursor.fetchall()]
            
            return AnalyticsData(
                total_connections=total_connections,
                total_commands=total_commands,
                success_rate=success_rate,
                top_commands=top_commands,
                enterprise_usage=enterprise_data,
                recent_activity=recent_activity
            )
    
    def get_enterprise_summary(self) -> List[Dict]:
        """Get enterprise connection summary"""
        with self.get_connection() as conn:
            cursor = conn.execute("""
                SELECT 
                    enterprise,
                    COUNT(*) as connection_count,
                    GROUP_CONCAT(DISTINCT tags) as all_tags,
                    MAX(last_used) as last_activity,
                    AVG(usage_count) as avg_usage
                FROM connections 
                WHERE enterprise IS NOT NULL
                GROUP BY enterprise
                ORDER BY connection_count DESC
            """)
            
            enterprises = []
            for row in cursor.fetchall():
                enterprise = dict(row)
                # Get common commands for this enterprise
                cmd_cursor = conn.execute("""
                    SELECT ch.command, COUNT(*) as count
                    FROM command_history ch
                    JOIN connections c ON ch.connection_id = c.id
                    WHERE c.enterprise = ?
                    GROUP BY ch.command
                    ORDER BY count DESC
                    LIMIT 5
                """, (enterprise['enterprise'],))
                
                enterprise['common_commands'] = [row['command'] for row in cmd_cursor.fetchall()]
                enterprises.append(enterprise)
            
            return enterprises
    
    def cleanup_old_data(self, days: int = 90):
        """Clean up old data to maintain performance"""
        cutoff_date = datetime.now() - timedelta(days=days)
        
        with self.get_connection() as conn:
            # Clean old command history
            cursor = conn.execute("""
                DELETE FROM command_history 
                WHERE timestamp < ?
            """, (cutoff_date,))
            
            deleted_count = cursor.rowcount
            conn.commit()
            
            logger.info(f"Cleaned up {deleted_count} old command history records")
    
    def close(self):
        """Close database connections"""
        # SQLite connections are automatically closed by context manager
        logger.info("Database manager closed")